package shardmaster

import "net"
import "fmt"
import "net/rpc"
import "paxos"
import "sync"
import "os"
import "syscall"
import "encoding/gob"
import "math/rand"
import cryptoRand "crypto/rand"
import "math/big"
import "time"
//import "strings"
import "util"

type ShardMaster struct {
	mu         sync.Mutex
	l          net.Listener
	me         int
	dead       bool // for testing
	unreliable bool // for testing
	px         *paxos.Paxos

	configs []Config // indexed by config num
	insNum  int      // instance/instruction number

	logger *util.Logger
}

const (
	JOIN  = "JOIN"
	LEAVE = "LEAVE"
	MOVE  = "MOVE"
	QUERY = "QUERY"
)

type Op struct {
	// Your data here.
	UID       int64
	OpType    string
	GID       int64
	Servers   []string
	ShardNum  int
	ConfigNum int
}

//
// used to generate a UID for an operation
//
func nrand() int64 {
	max := big.NewInt(int64(1) << 62)
	bigx, _ := cryptoRand.Int(cryptoRand.Reader, max)
	x := bigx.Int64()
	return x
}

func (sm *ShardMaster) ExecuteOp(op *Op) {
	//fmt.Printf("ExecuteOp: %v\n", *op)
	//fmt.Printf("Server %v No.%v Op: ", sm.me, sm.insNum);

	if op.OpType == QUERY {
		//fmt.Printf("Query\n")
		return
	}

	NumConfigs := len(sm.configs)
	newConfig := Config{}
	newConfig.Num = NumConfigs
	newConfig.Shards = sm.configs[NumConfigs-1].Shards
	newConfig.Groups = map[int64][]string{}
	for GID, Servers := range sm.configs[NumConfigs-1].Groups {
		newConfig.Groups[GID] = Servers
	}

	if op.OpType == MOVE {
		//fmt.Printf("Move shard %v to group %v\n", op.ShardNum, op.GID)
		newConfig.Shards[op.ShardNum] = op.GID
	}

	if op.OpType == JOIN {
		//fmt.Printf("Join group %v\n", op.GID)
		sm.ExecuteJoin(op, &newConfig)
	}

	if op.OpType == LEAVE {
		//fmt.Printf("Leave group %v\n", op.GID)
		sm.ExecuteLeave(op, &newConfig)
	}

	//fmt.Printf("\tCfg #%v  %v\n", newConfig.Num, newConfig.Shards)
	sm.configs = append(sm.configs, newConfig)
}

func (sm *ShardMaster) ExecuteJoin(op *Op, config *Config) {
	config.Groups[op.GID] = op.Servers
	NumGroups := len(config.Groups)

	// handle special case for empty config
	if NumGroups == 1 {
		for i := 0; i < NShards; i++ {
			config.Shards[i] = op.GID
		}
		return
	}

	N := NShards / NumGroups     // each group has N or N+1 shards
	X := NShards % NumGroups     // number of groups that have N+1 shards
	Y := NumGroups - X           // number of groups that have N shards
	groupSize := map[int64]int{} // map for storing the size of each group
	for GID := range config.Groups {
		groupSize[GID] = 0
	}
	for i := 0; i < NShards; i++ {
		groupSize[config.Shards[i]]++
	}

	for {
		findLess := false
		findMore := false
		GidLess := int64(-1)
		GidMore := int64(-1)
		NCount := 0
		Nplus1Count := 0
		for GID, size := range groupSize {
			if findLess == false {
				if size < N || (size == N && NCount == Y) {
					findLess = true
					GidLess = GID
				}
				if size == N {
					NCount++
				}
			}
			if findMore == false {
				if size > N+1 || (size == N+1 && Nplus1Count == X) {
					findMore = true
					GidMore = GID
				}
				if size == N+1 {
					Nplus1Count++
				}
			}
		}

		if findMore != findLess {
			fmt.Printf("ERROR: cannot find moving pair!!\n")
			fmt.Printf("N=%d X=%d Y=%d\n", N, X, Y)
			fmt.Println(groupSize)
			break
		}

		if findMore == false {
			break
		}

		groupSize[GidMore]--
		groupSize[GidLess]++
		for i := 0; i < NShards; i++ {
			if config.Shards[i] == GidMore {
				config.Shards[i] = GidLess
				break
			}
		}
	}
}

func (sm *ShardMaster) ExecuteLeave(op *Op, config *Config) {
	NumGroups := len(config.Groups)

	_, exists := config.Groups[op.GID]
	if exists == false {
		fmt.Printf("WARNING: Group %d not in the configuration\n", op.GID)
		return
	} else {
		NumGroups--
	}
	
	if NumGroups == 0 {
		for i, _ := range config.Shards {
			config.Shards[i] = 0
		}
		delete(config.Groups, op.GID)
		return
	}

	N := NShards / NumGroups     // each group has N or N+1 shards
	X := NShards % NumGroups     // number of groups that have N+1 shards
	Y := NumGroups - X           // number of groups that have N shards
	groupSize := map[int64]int{} // map for storing the size of each group
	for GID := range config.Groups {
		groupSize[GID] = 0
	}
	for i := 0; i < NShards; i++ {
		groupSize[config.Shards[i]]++
	}

	for {
		findLess := false
		findMore := false
		GidLess := int64(-1)
		GidMore := int64(-1)
		NCount := 0
		Nplus1Count := 0
		for GID, size := range groupSize {
			if findLess == false && GID != op.GID {
				if size < N || (size == N && NCount == Y) {
					findLess = true
					GidLess = GID
				}
				if size == N {
					NCount++
				}
			}
			if findMore == false {
				if (size > 0 && GID == op.GID) ||
					size > N+1 || (size == N+1 && Nplus1Count == X) {
					findMore = true
					GidMore = GID
				}
				if size == N+1 && GID != op.GID {
					Nplus1Count++
				}
			}
		}

		if findMore != findLess {
			fmt.Printf("ERROR: cannot find moving pair!!\n")
			fmt.Printf("N=%d X=%d Y=%d\n", N, X, Y)
			fmt.Println(groupSize)
			break
		}

		if findMore == false {
			break
		}

		groupSize[GidMore]--
		groupSize[GidLess]++
		for i := 0; i < NShards; i++ {
			if config.Shards[i] == GidMore {
				config.Shards[i] = GidLess
				break
			}
		}
	}
	delete(config.Groups, op.GID)
}

func (sm *ShardMaster) Join(args *JoinArgs, reply *JoinReply) error {
	sm.logger.Logger.Printf("Join %v %v\n", args.GID, args.Servers)
	sm.mu.Lock()
	defer sm.mu.Unlock()

	joinOp := Op{}
	joinOp.UID = nrand()
	joinOp.OpType = JOIN
	joinOp.GID = args.GID
	joinOp.Servers = args.Servers

	for {
		sm.px.Start(sm.insNum, joinOp)

		to := 10 * time.Millisecond
		for {
			decided, res := sm.px.Status(sm.insNum)
			if decided {
				sm.insNum += 1
				op := res.(Op)
				sm.ExecuteOp(&op)
				sm.px.Done(sm.insNum - 1)
				if op.UID == joinOp.UID {
					if op.OpType != joinOp.OpType {
						fmt.Printf("ERROR: Different ops with the same uid!!!\n")
					}
					return nil
				}
				break
			}

			time.Sleep(to)
			if to < 10*time.Second {
				to *= 2
			}
		}
	}
	return nil
}
func (sm *ShardMaster) SetTag(tag string) {
	sm.px.SetTag(tag)
}
func (sm *ShardMaster) Leave(args *LeaveArgs, reply *LeaveReply) error {
	sm.logger.Logger.Printf("Leave %v\n", args.GID)

	sm.mu.Lock()
	defer sm.mu.Unlock()

	leaveOp := Op{}
	leaveOp.UID = nrand()
	leaveOp.OpType = LEAVE
	leaveOp.GID = args.GID

	for {
		sm.px.Start(sm.insNum, leaveOp)

		to := 10 * time.Millisecond
		for {
			decided, res := sm.px.Status(sm.insNum)
			if decided {
				sm.insNum += 1
				op := res.(Op)
				sm.ExecuteOp(&op)
				sm.px.Done(sm.insNum - 1)
				if op.UID == leaveOp.UID {
					if op.OpType != leaveOp.OpType {
						fmt.Printf("ERROR: Different ops with the same uid!!!\n")
					}
					return nil
				}
				break
			}

			time.Sleep(to)
			if to < 10*time.Second {
				to *= 2
			}
		}
	}
	return nil

}

func (sm *ShardMaster) Move(args *MoveArgs, reply *MoveReply) error {

	sm.mu.Lock()
	defer sm.mu.Unlock()

	moveOp := Op{}
	moveOp.UID = nrand()
	moveOp.OpType = MOVE
	moveOp.GID = args.GID
	moveOp.ShardNum = args.Shard

	for {
		sm.px.Start(sm.insNum, moveOp)

		to := 10 * time.Millisecond
		for {
			decided, res := sm.px.Status(sm.insNum)
			if decided {
				sm.insNum += 1
				op := res.(Op)
				sm.ExecuteOp(&op)
				sm.px.Done(sm.insNum - 1)
				if op.UID == moveOp.UID {
					if op.OpType != moveOp.OpType {
						fmt.Printf("ERROR: Different ops with the same uid!!!\n")
					}
					return nil
				}
				break
			}

			time.Sleep(to)
			if to < 10*time.Second {
				to *= 2
			}
		}
	}
	return nil
}

func (sm *ShardMaster) Query(args *QueryArgs, reply *QueryReply) error {
	sm.logger.Logger.Printf("Query %v\n", args.Num)

	sm.mu.Lock()
	defer sm.mu.Unlock()

	queryOp := Op{}
	queryOp.UID = nrand()
	queryOp.OpType = QUERY
	queryOp.ConfigNum = args.Num

	for {
		sm.px.Start(sm.insNum, queryOp)

		to := 10 * time.Millisecond
		for {
			decided, res := sm.px.Status(sm.insNum)
			if decided {
				sm.insNum += 1
				op := res.(Op)
				sm.ExecuteOp(&op)
				sm.px.Done(sm.insNum - 1)
				if op.UID == queryOp.UID {
					if op.OpType != queryOp.OpType {
						fmt.Printf("ERROR: Different ops with the same uid!!!\n")
					}
					NumConfigs := len(sm.configs)
					if args.Num == -1 || args.Num >= NumConfigs {
						reply.Config = sm.configs[NumConfigs-1]
					} else {
						reply.Config = sm.configs[args.Num]
					}
					return nil
				}
				break
			}

			time.Sleep(to)
			if to < 10*time.Second {
				to *= 2
			}
		}
	}
	return nil

}
func (sm *ShardMaster) Alive() bool{
    return !sm.dead
}

func (sm *ShardMaster) CleanDir() {
    sm.px.CleanDir()
}
// please don't change this function.
func (sm *ShardMaster) Kill() {
	sm.dead = true
	sm.l.Close()
	sm.px.Kill()
    // some r/w may not be finished
    //sm.px.CleanDir()
}

//
// servers[] contains the ports of the set of
// servers that will cooperate via Paxos to
// form the fault-tolerant shardmaster service.
// me is the index of the current server in servers[].
// 
func StartServer(servers []string, me int) *ShardMaster {
	gob.Register(Op{})

	sm := new(ShardMaster)
	sm.me = me

	sm.configs = make([]Config, 1)
	sm.configs[0].Groups = map[int64][]string{}

	rpcs := rpc.NewServer()
	rpcs.Register(sm)
	sm.px = paxos.Make(servers, me, rpcs)

	sm.logger = util.NewLogger(servers[me] + ".sm", false)

	var l net.Listener
	var e error

	port := servers[me]
	if paxos.IsIP(port) {
		sm.logger.Logger.Printf("Listening to TCP %v\n", port)
		l, e = net.Listen("tcp", port)
	} else {
		sm.logger.Logger.Printf("Listening to File %v\n", port)
		os.Remove(port)
		l, e = net.Listen("unix", port)
	}

	if e != nil {
		sm.logger.Logger.Fatal("listen error: ", e)
	}
	sm.l = l

	// please do not change any of the following code,
	// or do anything to subvert it.

	go func() {
		for sm.dead == false {
			conn, err := sm.l.Accept()
			if err == nil && sm.dead == false {
				if sm.unreliable && (rand.Int63()%1000) < 100 {
					// discard the request.
					conn.Close()
				} else if sm.unreliable && (rand.Int63()%1000) < 200 {
					// process the request but force discard of reply.
					c1 := conn.(*net.UnixConn)
					f, _ := c1.File()
					err := syscall.Shutdown(int(f.Fd()), syscall.SHUT_WR)
					if err != nil {
						fmt.Printf("shutdown: %v\n", err)
					}
					go rpcs.ServeConn(conn)
				} else {
					go rpcs.ServeConn(conn)
				}
			} else if err == nil {
				conn.Close()
			}
			if err != nil && sm.dead == false {
				fmt.Printf("ShardMaster(%v) accept: %v\n", me, err.Error())
				sm.Kill()
			}
		}
	}()

	return sm
}

/*
func (t Config) String() string {
	return fmt.Sprintf("{%v %v}", t.Num, t.Shards)
}
*/
